# app/ws_manager.py
import asyncio
from fastapi import WebSocket
from typing import List, Dict, Optional
from app.core.logger import get_logger

logger = get_logger(__name__)


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.lock = asyncio.Lock()

    async def connect_device(self, websocket: WebSocket, dev_id: str):
        """异步连接设备"""
        async with self.lock:
            self.active_connections[dev_id] = websocket
        # logger.info(f"设备 {dev_id} 已连接")

    async def disconnect_device(self, dev_id: str):
        """异步断开设备连接"""
        async with self.lock:
            if dev_id in self.active_connections:
                try:
                    await self.active_connections[dev_id].close()
                except:
                    # 如果关闭失败，强制断开
                    try:
                        self.active_connections[dev_id].close()
                    except:
                        pass
                finally:
                    del self.active_connections[dev_id]
        logger.info(f"设备 {dev_id} 已断开连接")

    async def send_to_device(self, dev_id: str, message: dict):
        """异步发送消息给设备"""
        async with self.lock:
            if dev_id in self.active_connections:
                try:
                    await self.active_connections[dev_id].send_json(message)
                except:
                    await self.disconnect_device(dev_id)
    async def get_device_connection(self, dev_id: str) -> Optional[WebSocket]:
        """获取设备WebSocket连接"""
        async with self.lock:
            return self.active_connections.get(dev_id)


# 全局设备连接管理器
ws_manager = ConnectionManager()


class FrontendConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
        self.lock = asyncio.Lock()

    async def connect_frontend(self, websocket: WebSocket):
        """异步连接前端"""
        await websocket.accept()
        async with self.lock:
            self.active_connections.append(websocket)
        logger.info("前端连接已建立")

    async def disconnect_frontend(self, websocket: WebSocket):
        """异步断开前端连接"""
        async with self.lock:
            if websocket in self.active_connections:
                self.active_connections.remove(websocket)
        logger.info("前端连接已断开")

    async def broadcast(self, message: dict):
        """异步广播消息给所有前端"""
        disconnected = []
        async with self.lock:
            for connection in self.active_connections:
                try:
                    await connection.send_json(message)
                except:
                    disconnected.append(connection)

        # 清理断开连接的客户端
        for connection in disconnected:
            await self.disconnect_frontend(connection)

class DeviceInfoConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str,List[WebSocket]] = {}
        self.command_senders: Dict[str, WebSocket] = {}  # 存储命令发送者 {dev_id: websocket}
        self.lock = asyncio.Lock()

    async def connect(self,websocket:WebSocket,dev_id:str):
        """连接设备信息客户端"""
        async with self.lock:
            if dev_id not in self.active_connections:
                self.active_connections[dev_id] = set()  # 使用集合避免重复
            if websocket not in self.active_connections[dev_id]:
                self.active_connections[dev_id].add(websocket)  # 添加连接
                logger.info(f"设备信息WebSocket连接已建立: {dev_id}")
            else:
                logger.warning(f"设备 {dev_id} 的信息客户端已存在，跳过重复连接")

    async def disconnect(self,websocket:WebSocket,dev_id:str):
        """断开设备信息客户端连接"""
        async with self.lock:
            if dev_id in self.active_connections:
                self.active_connections[dev_id].remove(websocket)
                if not self.active_connections[dev_id]:
                    del self.active_connections[dev_id]
        logger.info(f"设备 {dev_id} 的信息客户端已断开连接")

    async def broadcast_device_info(self, dev_id: str, device_info: dict):
        """广播设备信息给所有客户端"""
        disconnected = []
        """广播设备信息给所有订阅者"""
        async with self.lock:
            if dev_id in self.active_connections:
                sent_websockets = set()  # 记录已发送的连接
                for websocket in self.active_connections[dev_id]:
                    if websocket not in sent_websockets:  # 确保只发送一次
                        try:
                            await websocket.send_json(device_info)
                            sent_websockets.add(websocket)
                        except Exception as e:
                            logger.error(f"广播设备信息失败: {e}")
        # 清理断开的连接
        for connection, dev_id in disconnected:
            await self.disconnect(connection, dev_id)
            logger.info(f"设备 {dev_id} 的信息客户端已断开连接")

    async def send_command(self, dev_id: str, command: dict):
        """发送命令给设备信息客户端"""
        async with self.lock:
            if dev_id in self.active_connections:
                for websocket in self.active_connections[dev_id]:
                    try:
                        await websocket.send_json(command)
                        return True
                    except Exception as e:
                        logger.error(f"发送命令给设备 {dev_id} 失败: {e}")
        return False

    def set_command_sender(self, dev_id: str, websocket: WebSocket):
        """设置命令发送者"""
        self.command_senders[dev_id] = websocket

    def get_command_sender(self, dev_id: str) -> Optional[WebSocket]:
        """获取命令发送者"""
        return self.command_senders.get(dev_id)

    def remove_command_sender(self, dev_id: str, websocket: WebSocket):
        """移除命令发送者记录"""
        if dev_id in self.command_senders and self.command_senders[dev_id] == websocket:
            del self.command_senders[dev_id]




device_info_manager = DeviceInfoConnectionManager()
# 全局前端连接管理器
frontend_manager = FrontendConnectionManager()
